# 导入 PyTorch 深度学习框架
import torch
# 从 PyTorch 中导入神经网络模块
import torch.nn as nn
# 从 torchvision 库中导入预训练的 MobileNetV2 和 ShuffleNetV2_x0_5 模型
from torchvision.models import mobilenet_v2, shufflenet_v2_x0_5


class SwitchModel(nn.Module):
    def __init__(self, backbone="shufflenet", num_classes=2, pretrained=True):
        """
        初始化 SwitchModel 类

        Args:
            backbone (str, optional): 模型的骨干网络，可选值为 "mobilenet" 或 "shufflenet"，默认为 "shufflenet"
            num_classes (int, optional): 分类的类别数量，默认为 2
            pretrained (bool, optional): 是否使用预训练的模型权重，默认为 True
        """
        # 调用父类 nn.Module 的构造函数
        super().__init__()
        
        # 构建模型
        # 确保输入的骨干网络名称为 "mobilenet" 或 "shufflenet"
        assert backbone in ["mobilenet", "shufflenet"]
        if backbone == "mobilenet":
            # 使用预训练的 MobileNetV2 模型
            self.model = mobilenet_v2(pretrained=pretrained)
            # 假设使用 MobileNetV2 的最后一层特征维度
            last_channel = self.model.last_channel
            # 修改模型的分类器部分，添加 Dropout 层并调整全连接层的输出维度
            self.model.classifier[1] = nn.Sequential(
                    nn.Dropout(p=0.2),
                    nn.Linear(last_channel, num_classes),
                )
        else:
            # 使用预训练的 ShuffleNetV2_x0_5 模型
            self.model = shufflenet_v2_x0_5(pretrained=pretrained)
            # 修改模型的全连接层，调整输出维度
            self.model.fc = nn.Linear(1024, num_classes)

    def forward(self, x):
        """
        前向传播函数

        Args:
            x (torch.Tensor): 输入的图像张量

        Returns:
            torch.Tensor: 模型的输出张量
        """
        # 通过模型进行前向传播
        output = self.model(x)
        return output
